from mmseg.datasets.builder import DATASETS
from mmseg.datasets.custom import CustomDataset


@DATASETS.register_module()
class Camvid32(CustomDataset):
    """ADE20K dataset.

    In segmentation map annotation for ADE20K, 0 stands for background, which
    is not included in 150 categories. ``reduce_zero_label`` is fixed to True.
    The ``img_suffix`` is fixed to '.jpg' and ``seg_map_suffix`` is fixed to
    '.png'.
    """
    # CLASSES =('Bicyclist','Building','Car','Column_Pole','Fence','Pedestrian','Road'
    #     'Sidewalk','SignSymbol','Sky','Tree')
    # PALETTE = [[0, 128, 192],[128, 0, 0],[64, 0, 128],[192, 192, 128],[64, 64, 128],[64, 64, 0],[128, 64, 128],
    #    [0, 0, 192],[192, 128, 128], [128, 128, 128],[128, 128, 0]]

    CLASSES = (
        'Animal', 'Archway', 'Bicyclist', 'Bridge', 'Building', 'Car', 'CartLuggagePram', 'Child', 'Column_Pole',
        'Fence',
        'LaneMkgsDriv', 'LaneMkgsNonDriv', 'Misc_Text', 'MotorcycleScooter', 'OtherMoving', 'ParkingBlock',
        'Pedestrian',
        'Road', 'RoadShoulder', 'Sidewalk', 'SignSymbol', 'Sky', 'SUVPickupTruck', 'TrafficCone', 'TrafficLight',
        'Train',
        'Tree', 'Truck_Bus', 'Tunnel', 'VegetationMisc', 'Void', 'Wall')

    PALETTE = [[64, 128, 64], [192, 0, 128], [0, 128, 192], [0, 128, 64], [128, 0, 0], [64, 0, 128], [64, 0, 192],
               [192, 128, 64],
               [192, 192, 128], [64, 64, 128], [128, 0, 192], [192, 0, 64], [128, 128, 64], [192, 0, 192],
               [128, 64, 64],
               [64, 192, 128], [64, 64, 0], [128, 64, 128], [128, 128, 192], [0, 0, 192], [192, 128, 128],
               [128, 128, 128],
               [64, 128, 192], [0, 0, 64], [0, 64, 64], [192, 64, 128], [128, 128, 0], [192, 128, 192], [64, 0, 64],
               [192, 192, 0], [0, 0, 0], [64, 192, 0]]

    def __init__(self, **kwargs):
        super(Camvid32, self).__init__(
            img_suffix='.png',
            seg_map_suffix='.png',
            reduce_zero_label=True,
            **kwargs)
